from scipy.misc import imread
from scipy.misc import imsave
import numpy as np
import logging
from logging import DEBUG as debug
from logging import log
import argparse

from numpy import empty, arange, exp, real, imag, pi
from numpy.fft import rfft, irfft
import pywt
from math import floor


######################################################################
# 1D DCT Type-II
def dct(y):
    N = len(y)
    y2 = empty(2 * N, float)
    y2[:N] = y[:]
    y2[N:] = y[::-1]

    c = rfft(y2)
    phi = exp(-1j * pi * arange(N) / (2 * N))
    return real(phi * c[:N])


######################################################################
# 1D inverse DCT Type-II
def idct(a):
    N = len(a)
    c = empty(N + 1, complex)

    phi = exp(1j * pi * arange(N) / (2 * N))
    c[:N] = phi * a
    c[N] = 0.0
    return irfft(c)[:N]


######################################################################
# 2D DCT
def dct2(y):
    M = y.shape[0]
    N = y.shape[1]
    a = empty([M, N], float)
    b = empty([M, N], float)

    for i in range(M):
        a[i, :] = dct(y[i, :])
    for j in range(N):
        b[:, j] = dct(a[:, j])

    return b


######################################################################
# 2D inverse DCT
def idct2(b):
    M = b.shape[0]
    N = b.shape[1]
    a = empty([M, N], float)
    y = empty([M, N], float)

    for i in range(M):
        a[i, :] = idct(b[i, :])
    for j in range(N):
        y[:, j] = idct(a[:, j])

    return y


######################################################################
# 1D DST Type-I
def dst(y):
    N = len(y)
    y2 = empty(2 * N, float)
    y2[0] = y2[N] = 0.0
    y2[1:N] = y[1:]
    y2[:N:-1] = -y[1:]
    a = -imag(rfft(y2))[:N]
    a[0] = 0.0

    return a


######################################################################
# 1D inverse DST Type-I
def idst(a):
    N = len(a)
    c = empty(N + 1, complex)
    c[0] = c[N] = 0.0
    c[1:N] = -1j * a[1:]
    y = irfft(c)[:N]
    y[0] = 0.0

    return y


######################################################################
# 2D DST
def dst2(y):
    M = y.shape[0]
    N = y.shape[1]
    a = empty([M, N], float)
    b = empty([M, N], float)

    for i in range(M):
        a[i, :] = dst(y[i, :])
    for j in range(N):
        b[:, j] = dst(a[:, j])

    return b


######################################################################
# 2D inverse DST
def idst2(b):
    M = b.shape[0]
    N = b.shape[1]
    a = empty([M, N], float)
    y = empty([M, N], float)

    for i in range(M):
        a[i, :] = idst(b[i, :])
    for j in range(N):
        y[:, j] = idst(a[:, j])

    return y


def main():
    # Part I
    FORMAT = '	[%(levelname)-5s]%(asctime)-8s %(filename)s:%(lineno)d %(message)s'
    DATEFORMAT = '%H:%M:%S'
    logging.basicConfig(level=logging.DEBUG, format=FORMAT, datefmt=DATEFORMAT)
    logging.log(logging.DEBUG, "Start")
    parser = argparse.ArgumentParser()
    parser.add_argument("--infile", help="Input Image File", type=str, default='data/lenna.tif')
    parser.add_argument("--pitchSize", help="Input Image File", type=int, default=8)
    args = parser.parse_args()

    inFileName = args.infile
    patchSize = args.pitchSize
    log(debug, "Handle file " + inFileName)
    log(debug, "Patch size " + str(patchSize))

    originalImage = imread(inFileName)

    pattern = np.zeros((patchSize, patchSize))
    for row in range(patchSize):
        for col in range(patchSize):
            pattern[row, col] = row ** 2 + col ** 2
    pattern = pattern < 10
    pattern = np.double(pattern)

    resultImage = np.zeros(originalImage.shape)
    for row in range(0, originalImage.shape[0], 8):
        for col in range(0, originalImage.shape[1], 8):
            pitch = originalImage[row:row + 8, col:col + 8]
            afterdct = dct2(pitch)
            afterdct = afterdct  * pattern
            afteridcr = idct2(afterdct)
            resultImage[row:row+8, col:col+8] = afteridcr

    imsave('out/hw7_dct_result.jpg', resultImage)
    imsave('out/hw7_dct_diff.jpg', np.abs(resultImage - originalImage))
    imsave('out/hw7_zoriginal.jpg', originalImage)

    # Part II
    for name in ['haar', 'db8','sym8','bior6.8']:
        coeffs = pywt.wavedec2(originalImage, 'haar', level=3)
        arr, coeff_slices = pywt.coeffs_to_array(coeffs)
        arr = np.array(arr)
        num = np.abs(arr)
        num = np.reshape(num,(num.shape[0]*num.shape[1],1))
        threshHole = num[floor(num.shape[0] * 0.4)]
        arr[np.abs(arr) < threshHole] = 0
        coeffs = pywt.array_to_coeffs(arr,coeff_slices, 'wavedec2')
        restored = pywt.waverec2(coeffs,'haar')
        imsave('out/hw7_wavelet_%s_result.jpg' % name, restored)
        imsave('out/hw7_wavelet_%s_diff.jpg' % name, np.abs(restored - originalImage))
if __name__ == '__main__':
    main()
